import os
import re
import uuid

import click
import dotenv
import pandas as pd
import requests
from tqdm import tqdm


@click.command()
@click.argument("file", type=click.Path(exists=True))
@click.option("--dataset-id", help="dataset id", required=True)
@click.option("--document-id", help="document id", required=True)
@click.option("--sheet-name", help="sheet name", default=None)
def main(file, dataset_id, document_id, sheet_name):
    """Upload dataset to Dify from Excel file."""
    dotenv.load_dotenv()
    api_key = os.getenv("DIFY_API_KEY_DB")
    base_url = os.getenv("DIFY_BASE_URL")

    if not sheet_name:
        df = pd.read_excel(file)
    else:
        df = pd.read_excel(file, sheet_name=sheet_name)

    for _, row in tqdm(df.iterrows()):
        url = f"{base_url}/v1/datasets/{dataset_id}/documents/{document_id}/segments"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        tags = row["tag"]
        if not isinstance(tags, str):
            keywords = []
        else:
            tags = [line.strip() for line in tags.split("\n") if line.strip()]
            keywords = []

            for tag in tags:
                tag = re.sub(r"一级事项引导[：:]", "", tag)
                tag = re.sub(r"二级事项引导[：:]", "", tag)
                tag = re.sub(r"场景引导[：:]", "", tag)
                tag = re.sub(r"关键词[：:]", "", tag)
                keywords.extend(re.split(r"[,，]", tag))

        data = {
            "segments": [
                {
                    "content": f"{'-'.join(keywords)}\n{row['question']}\n{row['answer']}",
                    "answer": row["answer"],
                    "keywords": keywords,
                }
            ]
        }
        try:
            str(uuid.UUID(row["id"]))
        except:
            row["id"] = str(row["id"])
            row["id"] = row["id"] + "0" * (32 - len(row["id"]))

        if "id" in row and not pd.isna(row["id"]):
            cur_url = os.path.join(url, str(uuid.UUID(row["id"])))
            response = requests.delete(cur_url, headers=headers)

            data["segments"][0]["segment_id"] = str(uuid.UUID(row["id"]))
            response = requests.post(url, headers=headers, json=data)
        else:
            response = requests.post(url, headers=headers, json=data)


if __name__ == "__main__":
    main()
